from typing import Optional
from matplotlib import pyplot as plt
import numpy as np
import math
import copy
from cion.gates import BaseGate, AWGBaseGate, AdvanceGate, sync, awg_trigger, pending

class Sequence:
    def __init__(self, ion_number = 5, chapter_dict = None, sampling_rate = 1000):
        self.awg_sampling_rate = 1000 #unit: MHz /default: 625 MHz
        self.awg_trigger = False
        self.awg_trigger_time = -1 #awg_trigger打开的时间
        self.current_time = -1 #表示上一个门结束的时间
        self.input_list = []
        self.set_para_flag = False
        self.gate_sequence = []
        self.ion_number = ion_number # total ions
        self.time_stamp = [0 for i in range(self.ion_number)]
        #self.gate_list = []
        self.gate_list_per_ion = [[] for i in range(ion_number)]
        self.labelled_pulse = {}
        self.awg_data = None
        self.real_chapter_dict = chapter_dict

    # set_parameters用于修改label对应的那些pulse的参数
    #请注意，如果想对一组pulse赋予相同的值，请设置batch_mode = True。这时程序会无条件将相同的参数赋值给每一个pulse
    #例如：set_parameters(label='rx0', duration=100, amp = 10, batch_mode = True)
    #表示将所有label为'rx0'的门的duration设为100，amp设为10。即批量赋值
    #若是set_parameters(label='rx0', duration=100)或set_parameters(label='rx0', duration=100, batch_mode = False)
    #且label为'rx0'的门至少有两个时，程序会报错
    ########################
    ########################
    #如果想同时修改多个pulse的某个参数，并且对每个pulse赋予不同的值，请一定使用tuple类型！！使用list会产生歧义。
    ########################
    #例如：有三个pulse的label都是'Rx'时，调用如下函数
    # amp = [(amp1, amp2, amp3) for i in range(12)]
    #set_parameters(label = 'Rx', duration = duration, amp = [amp1, amp2, amp3], freq = (freq1,freq2,freq3))
    #这里，程序会报错，因为并没有设置batch_mode = True，程序发现amp(虽然是list)是赋值给单个pulse的，无法分配给三个pulse
    #而如下参数形式是允许的：
    #set_parameters(label = 'Rx', duration = duration, freq = (freq1,None,freq3), phase = (phase0, [phase1,phase2,phase3], phase4))
    #这表明第二个pulse的phase参数确实是一个list，并且被修改为[phase1,phase2,phase3]，list用于处理分段函数的情况
    #而phase1,phase2,phase3本身也可以是tuple类型，用于处理每一段函数内部有多个sine分量相加的情况，具体请看相关示例
    ########################
    #而使用tuple类型分开赋值时，请确保tuple的元素个数与需要修改的pulse个数相等
    #不想修改的直接传None——如上面freq = (freq1,None,freq3)表示第2个pulse的freq参数不做修改
    def set_parameters(self, label = None, duration = None, amp = None, freq = None, phase = None, para_table = None, batch_mode = False):
        assert label != None

        if self.labelled_pulse.get(label) == None:
            assert False
        self.set_para_flag = True

        gate_category = -1
        for gate in self.labelled_pulse[label]:
            if isinstance(gate, BaseGate):
                assert gate_category in {-1,1}
                gate_category = 1
            else:
                assert gate_category in {-1,2}
                gate_category = 2


        if gate_category == 1:
            all_paras = [duration, amp, freq, phase]
            num = len(self.labelled_pulse[label])
            para_names = ['duration', 'amp', 'freq', 'phase']
            if num > 1:
                if batch_mode == False:
                    for idx,para in enumerate(all_paras):
                        if para != None:
                            try:
                                assert type(para) == tuple
                                assert len(para) == num
                            except Exception as e:
                                print("Error during setting parameters! The shape of \'{}\' does not match the number of pulses {}".format(para_names[i], num))
                                raise e

            for i in range(4):
                para = all_paras[i]
                if para != None:
                    if batch_mode == False:
                        if type(para) != tuple:
                            para = [para for j in range(num)]
                        assert len(para) == num
                        for j in range(num):
                            if para[j] != None:
                                if i == 0:
                                    self.labelled_pulse[label][j].duration = copy.copy(para[j])
                                if i == 1:
                                    self.labelled_pulse[label][j].amp = copy.copy(para[j])
                                if i == 2:
                                    self.labelled_pulse[label][j].freq = copy.copy(para[j])
                                if i == 3:
                                    self.labelled_pulse[label][j].phase = copy.copy(para[j])
                    else: # batch_mode : allocate identical parameters to all pulses
                        for j in range(num):
                                if i == 0:
                                    self.labelled_pulse[label][j].duration = copy.copy(para)
                                if i == 1:
                                    self.labelled_pulse[label][j].amp = copy.copy(para)
                                if i == 2:
                                    self.labelled_pulse[label][j].freq = copy.copy(para)
                                if i == 3:
                                    self.labelled_pulse[label][j].phase = copy.copy(para)
        else:
            assert para_table != None
            if batch_mode == False:
                if type(para_table) !=  tuple:
                    para_table = [para_table]
                assert len(para_table) == len(self.labelled_pulse[label])
                for idx, single_dict in para_table:
                    if single_dict != None:
                        self.labelled_pulse[label][idx].set_parameter_table( copy.copy(single_dict) )
            else:
                for gate in self.labelled_pulse[label]:
                    gate.set_parameter_table( copy.copy(para_table) )
        '''
        if duration != None:
            if type(duration) != tuple:
                duration = [duration for i in range(num)]
            assert len(duration) == num
            for i in range(num):
                if duration[i] != None:
                    self.labelled_pulse[label][i].duration = duration[i]
        '''

    #continuous_seq = [RSB(paras,10.0),opitical_pumping(102,13),...]

    #seq.add_sequence(EIT(100...),continuous_seq,optical_pumping)

    def __add_gate_non_composite(self,gate):
        #if isinstance(gate, tuple):
        #    print(gate)
        ion = gate.ion_index

        assert ion != None
        tn = gate.duration
        tl = gate.latency
        if gate.label != None:
            if self.labelled_pulse.get(gate.label) != None:
                self.labelled_pulse[gate.label].append(gate)
            else:
                self.labelled_pulse[gate.label] = [gate]

        if isinstance(ion, int):
            self.gate_sequence.append((ion, self.time_stamp[ion]+tl, tn, gate.pulse_type, gate))
            self.time_stamp[ion] += tn + tl
            self.current_time = self.time_stamp[ion]
            if gate.awg_flag == True:
                assert self.awg_trigger == True
                self.gate_list_per_ion[ion].append((self.time_stamp[ion]-tn, tn, gate))
        else:
            tmax = -1
            for single_ion in ion:
                tmax = max(tmax, self.time_stamp[single_ion])
            self.gate_sequence.append((ion, tmax+tl, tn, gate.pulse_type, gate))
            for single_ion in ion:
                self.time_stamp[single_ion] = tmax + tn +tl
                if gate.awg_flag == True:
                    self.gate_list_per_ion[single_ion].append((self.time_stamp[single_ion]-tn, tn, gate))
            self.current_time = tmax + tn +tl

    def __add_gate(self, gate):
        if gate.pulse_type == None:
            print("Caution! You are defining pulses with no type! The the is below:")
            print(gate)

        # composite_flag: 表示判断此pulse是否有head或tail，即主pulse前面或者后面有一段或多段附加的pulse（通常时间很短）。
        ## 例如：某个DDS上的pulse持续100us（微秒），但是实验者可能会在其前面加1段持续1us的特殊pulse，在其后面加2段1us的特殊pulse。
        ## 此时，这个DDS pulse实际总共耗时103us。而我们的cion会将其处理成4个独立的pulse，但是这4个pulse一定是同时且连续出现的。
        # 当head和tail至少存在一个时，composite_flag为True。由硬件相关的fpga chapter的规则来定义。
        composite_flag = False
        composite_flag, single_gate_chapter = self.whether_composite_according_to_fpga_chapter(gate)
        if composite_flag:
            self.__add_gate_composite(gate)
        else:
            assert  single_gate_chapter != None
            gate.chapter_data = single_gate_chapter
            self.__add_gate_non_composite(gate)


    def __add_parallel_gate(self, ion_index, advanced_gate):
        composite_flag = False
        gate_chapter = self.real_chapter_dict[advanced_gate.pulse_type]
        if advanced_gate.awg_flag == False and len(gate_chapter) > 1:
            composite_flag = True
        if composite_flag:
            #TODO
            raise Exception('composite advanced gate currently not permitted')
            pass
        else:
            advanced_gate.chapter_data = self.real_chapter_dict[advanced_gate.pulse_type][0]
            self.__add_parallel_gate_nocomposite(ion_index, advanced_gate)


    def __add_parallel_gate_nocomposite(self, ion_index, advanced_gate):
        assert ion_index != None
        tn = advanced_gate.duration
        tl = advanced_gate.latency
        tmax = -1
        for single_ion in iter_int_or_tuple(ion_index):
            tmax = max(tmax, self.time_stamp[single_ion])
        self.gate_sequence.append((ion_index, tmax+tl, tn, advanced_gate.pulse_type, advanced_gate))
        for single_ion in iter_int_or_tuple(ion_index):
            self.time_stamp[single_ion] = tmax + tn +tl
            if advanced_gate.awg_flag == True:
                duration = advanced_gate.duration
                latency = advanced_gate.latency
                segment_number = advanced_gate.para_table['segment_number']
                time_intervals = advanced_gate.para_table['time_intervals']

                ion_data = advanced_gate.para_table['data_per_ion'][single_ion]
                amp = ion_data['amp']
                freq = ion_data['freq']
                phase = ion_data['phase']
                single_qubit_gate = AWGBaseGate(duration, latency=latency, amp=amp, freq=freq, phase=phase,
                                             awg_flag=True, segment_number=segment_number, time_intervals = time_intervals)
                self.gate_list_per_ion[single_ion].append((self.time_stamp[single_ion]-tn, tn, single_qubit_gate))
        self.current_time = tmax + tn +tl

    def add_gates(self, *args):
        self.input_list.extend(list(args))
        self.inner_add_gates(args)

    def inner_add_gates(self, args):
        gate_num = len(args)
        for i in range(gate_num):
            if isinstance(args[i], BaseGate):
                self.__add_gate(args[i])
            elif isinstance(args[i], list):
                sub_sequence = args[i]
                for pulse in sub_sequence:
                    if isinstance(pulse, BaseGate):
                        self.__add_gate(pulse)
                    elif isinstance(pulse, tuple):
                        tmax = -1
                        if pulse[0] == 'sync':
                            for single_ion in pulse[1]:
                                tmax = max(tmax, self.time_stamp[single_ion])
                            for single_ion in pulse[1]:
                                self.time_stamp[single_ion] = tmax
                            self.current_time = tmax
                        elif pulse[0] == 'awg_trigger':
                            assert self.awg_trigger == False
                            self.awg_trigger = True
                            self.awg_trigger_time = self.current_time + pulse[1]
                        elif pulse[0] == 'advanced_gate':
                            #print(pulse)
                            self.__add_parallel_gate(pulse[1], pulse[2])
                        else:
                            print('Error in pulse sequence, exception type not permitted')
                            raise Exception()
                        #self.gate_list.append()
            elif isinstance(args[i], tuple):
                pulse = args[i]
                tmax = -1
                if pulse[0] == 'sync':
                    for single_ion in pulse[1]:
                        tmax = max(tmax, self.time_stamp[single_ion])
                    for single_ion in pulse[1]:
                        self.time_stamp[single_ion] = tmax
                    self.current_time = tmax
                elif pulse[0] == 'awg_trigger':
                        self.awg_trigger = True
                        self.awg_trigger_time = self.current_time + pulse[1]
                elif pulse[0] == 'advanced gate':
                    #print(pulse)
                    self.__add_parallel_gate(pulse[1],pulse[2])
                else:
                    pass

    def set_sequence(self, *args):
        self.clear(make_sure = True)
        self.input_list = list(args)
        self.inner_add_gates(args)

    def print_info(self):
        tmax = 0
        for pulse in self.gate_sequence:
            tmax = max(tmax, pulse[1]+pulse[2])
        print('----sequence info----')
        print('total pulses: {}'.format(len(self.gate_sequence)))
        print('sequence time: {} microseconds'.format(tmax))
        #print(self.gate_sequence)
        print('---------------------')


    def reformat_sequence(self):
        '''
        start_point : Dict; each element is a list of 3-tuples: (sequence_index, gate_object, ion_index)
        end_point : Dict; similar to start_point
        '''
        seq = self.gate_sequence
        time_index = set()

        for pulse in seq:
            time_index.add(pulse[1])
            time_index.add(pulse[1]+pulse[2])
        time_index = list(time_index)
        time_index.sort()
        #print(time_index)
        start_point = {}
        end_point = {}
        for idx,pulse in enumerate(seq):
            lt = pulse[1]
            rt = pulse[1]+pulse[2]
            if start_point.get(lt) != None:
                start_point[lt].append((idx,pulse[3],pulse[0], pulse[4]))
            else:
                start_point[lt] = [(idx,pulse[3],pulse[0], pulse[4])]
            if end_point.get(rt) != None:
                end_point[rt].append((idx,pulse[3],pulse[0], pulse[4]))
            else:
                end_point[rt] = [(idx,pulse[3],pulse[0], pulse[4])]

        #print(start_point, end_point)
        return time_index, start_point, end_point

    def get_chapter_data(self, pulse_type, ion_index):
        '''
        pulse_type is a string, such as 'EIT', 'optical_pumping', 'RSB' ...
        ion_index is an int or tuple, such as: 0,2,3 or (3), (0,2), (1,3,4)
        '''
        #this function is left for modification to get real chapter data according to the experiment configurations.
        return self.real_chapter_dict[pulse_type]

    def callable_test(self,func,t):
        if callable(func):
            result = func(t)
        else:
            result = func
        return result

    def func_generator(self,args):
        if isinstance(args['amp'],tuple):
            func_list_temp = []
            for i in range(len(args['amp'])):
                def func_temp(t):
                    return self.callable_test(args['amp'][i], t)*np.sin(self.callable_test(args['freq'][i], t)*t+self.callable_test(args['phase'][i], t))
                func_list_temp.append(func_temp)
            def func(t):
                result = 0
                for i in range(len(func_list_temp)):
                    result = result + func_list_temp[i](t)
                return result
        else:
            def func(t):
                try:
                    return self.callable_test(args['amp'], t)*np.sin(self.callable_test(args['freq'], t)*t+self.callable_test(args['phase'], t))
                except TypeError:
                    print(args['amp'], args['freq'], args['phase'])
                    raise TypeError
                    #pass

        return func

    def func_list_generator(self):
        #对于gate不分段的情况：
        #input is self.parameter_list = [{'amp':(amp_11,amp_12,...),'freq': (freq_11,freq_12,...), 'phase': (phase_11, phase_12,...)}, {'amp': amp_2, 'freq': freq_2, 'phase': phase_2}...]
        #如果gate本身是分段的，这里需要做一下处理：
        #input is self.parameter_list = [{'amp':[amp_seg1, amp_seg2, ...],'freq': [freq_seg1,freq_seg2,...], 'phase': [phase_seg1, ...]},...]
        #其中amp_seg1又是(amp_11,amp_12,...)的形式, 对应多个sin分量
        func_list = []
        t_list = []
        for gate_list in self.gate_list_per_ion:
            func_list_single_ion = []
            t_list_single_ion = []

            #print(func_list_single_ion)

            for gate in gate_list:
                #print(gate[0],gate[1],gate[2])
                if type(gate[2].amp) == list and len(gate[2].amp) == gate[2].segment_number:
                    for idx in range(len(gate[2].time_intervals)):
                        args = gate[2].parameter_list(idx)
                        try:
                            assert args['amp'] != None
                            assert args['freq'] != None
                            assert args['phase'] != None
                        except AssertionError as e:
                            print('Error in generating function, some parameters of a gate on ion {} are None.'.format(gate[2].ion_index))
                            raise e
                        lt, rt = gate[2].time_intervals[idx]
                        if rt <= gate[2].duration:
                            t_list_single_ion.append((gate[0]+lt, gate[0]+rt))
                            func_list_single_ion = func_list_single_ion + [self.func_generator(gate[2].parameter_list(idx))]
                        elif lt <= gate[2].duration:
                            t_list_single_ion.append((gate[0]+lt, gate[0]+gate[2].duration))
                            func_list_single_ion = func_list_single_ion + [self.func_generator(gate[2].parameter_list(idx))]
                else:
                    args = gate[2].parameter_list()
                    try:
                        assert args['amp'] != None
                        assert args['freq'] != None
                        assert args['phase'] != None
                    except AssertionError as e:
                        print('Error in generating function, some parameters of a gate on ion {} are None.'.format(gate[2].ion_index))
                        raise e
                    func_list_single_ion = func_list_single_ion + [self.func_generator(gate[2].parameter_list())]
                    t_list_single_ion.append((gate[0], gate[0]+gate[1]))

            func_list.append(func_list_single_ion)
            t_list.append(t_list_single_ion)

        return (func_list, t_list)

    # 此函数用于将每个离子上用的awg波形序列逐点计算出来，由于每个gate在awg上用的都是正弦波(sine function)，所以此函数对于不同实验室应该是通用的
    # Cion允许gate在awg上的函数是分段函数，并且每一段都可以是多个正弦分量的叠加
    def generator_awg_waveform(self):
        func_list, t_list = self.func_list_generator()
        data_list = []
        ion_id = 0
        for ion_f_list,ion_t_list in zip(func_list,t_list):
            tlist_temp = []
            flist_temp = np.array([])
            for i in range(len(ion_t_list)):
                lt, rt = ion_t_list[i]
                #t_temp = np.arange(lt, rt, 0.001)
                t_temp = np.arange(lt, rt, 1/self.awg_sampling_rate)
                t_temp = t_temp
                #print('lt,rt,t_temp:', lt,rt,len(t_temp))
                tlist_temp.append((lt, rt))
                if i > 0:
                    if lt > tlist_temp[-2][1]:
                        #num_zeroes = int((lt - tlist_temp[-2][1])*1000)
                        num_zeroes = round((lt - tlist_temp[-2][1])*self.awg_sampling_rate)
                        if num_zeroes % 10 != 0:
                            print(tlist_temp[-2][1],lt,num_zeroes)
                        zero_padd = np.zeros(num_zeroes)
                        flist_temp = np.append(flist_temp,zero_padd)
                elif (lt > 0): #i==0
                    num_zeroes = int(lt*1000)
                    zero_padd = np.zeros(num_zeroes)
                    flist_temp = np.append(flist_temp,zero_padd)

                try:
                    flist_temp = np.append(flist_temp, ion_f_list[i](t_temp))
                    #print(len(flist_temp))
                except TypeError as e:
                    print(lt, rt)
                    print(e)
                    raise e
                #start = ion_t_list[i+1]
            flist_temp = np.array(flist_temp)
            data_list.append((tlist_temp, flist_temp))
            ion_id += 1
        return data_list

    def regenerate_sequence(self):
        args = tuple(self.input_list)
        self.clear(make_sure = True)
        self.inner_add_gates(args)
        self.set_para_flag = False


##### 与硬件相关的接口函数
    # update_hardwares: 更新当前序列下label指定的这些gates(pulses)的硬件参数
    # reset_hardwares: 重置当前序列下label指定的这些gates的硬件参数
    # whether_composite_according_to_fpga_chapter: 通过fpga chapter来判断一个gate是否有head和tail
    # __add_gate_composite: FPGA编码规则下，处理带head或tail的组合gate，将其拆成多段来处理
    # perform_sequence_awg: 生成awg硬件所需要的长波形数据
    # perform_sequence_serial: 根据FPGA chapter_dict以及gates的时序关系，生成一个序列对应在FPGA所需要的最终数据

    # 每种类型的gate的硬件接口由用户决定：
    def update_hardwares(self, label=None):
        if isinstance(label, str):
            self.labelled_pulse[label][0].update_hardware()
        else:
            for i in label:
                self.labelled_pulse[i][0].update_hardware()

    def reset_hardwares(self, label=None):
        if isinstance(label, str):
            self.labelled_pulse[label][0].reset_hardware()
        else:
            for i in label:
                self.labelled_pulse[i][0].reset_hardware()

    def whether_composite_according_to_fpga_chapter(self, gate):
        gate_chapter = self.real_chapter_dict[gate.pulse_type]
        # awg上的gate已经定死，不能加head或tail
        if gate.awg_flag == False and len(gate_chapter) > 1:
            return True, None
        else:
            return False, gate_chapter[0]

    def __add_gate_composite(self, gate):
        chapter_list = self.real_chapter_dict[gate.pulse_type]###这一行涉及到具体的FPGA chapter，在不同硬件环境下不通用。后续升级需要将其调整至硬件相关的函数
        ion_index = gate.ion_index
        for chapter in chapter_list:
            if isinstance(chapter, str):
                #这里是主pulse
                gate.chapter_data = chapter
                self.__add_gate_non_composite(gate)
            else:
                #这里是head和tail，cion会将每一段头和尾都当成一个新的gate（BaseGate类型）
                t = chapter[1]
                temp_gate = BaseGate(duration=t, awg_flag=False)
                temp_gate.chapter_data = chapter[0] ###这一行涉及到具体的FPGA chapter，在不同硬件环境下不通用。后续升级需要将其调整至硬件相关的函数
                self.__add_gate_non_composite( temp_gate.on(ion_index) )

    def perform_sequence_awg(self):
        # generate the awg waveform data according to pulse waveform function
        _temp_awg_data = self.generator_awg_waveform()

        max_awg_data_len = 0
        for i in range(len(_temp_awg_data)):
            if len(_temp_awg_data[i][1]) > max_awg_data_len:
                max_awg_data_len = len(_temp_awg_data[i][1])
        self.awg_data = []
        for i in range(len(_temp_awg_data)):
            temp_data = np.pad(_temp_awg_data[i][1], (0,max_awg_data_len-len(_temp_awg_data[i][1])))
            self.awg_data.append(temp_data)
        self.awg_data = np.array(self.awg_data) ## This is what we need to send to the awg

    def perform_sequence_serial(self):

        # generate chapter data according to the datasheet and invoke the serial module
        # need the explict rules of parallel pulses

        if self.set_para_flag == True:
            print("Caution! You have modified some parameters, but no sequence updating executed. If needed, try use the member function \
\'regenerate_sequence()\' to update it. This run will perform the sequence without the parameters updated")

        time_index, start_point, end_point = self.reformat_sequence()
        active_set = set()
        ll = len(time_index)
        result_list = []
        #print(start_point)
        #print(end_point)
        #使用一个扫描队列来获取pulses的时序，从而得知每个时间段内有哪几个pulses是激活（正在作用的）的
        for idx, t in enumerate(time_index):
            if end_point.get(t) != None:
                for item in end_point.get(t):
                    # item: (sequence_index, pulse_type, ion_index, gate_object)
                    #active_set.remove(item[1])
                    if isinstance(item[2], int):
                        active_set.remove((item[1], item[2], item[3]))
                    else:
                        active_set.remove((item[1], tuple(item[2]),  item[3]))
            if start_point.get(t) != None:
                for item in start_point.get(t):
                    if isinstance(item[2], int):
                        active_set.add((item[1], item[2], item[3]))
                    else:
                        active_set.add((item[1], tuple(item[2]), item[3]))
            if idx < ll-1:
                chapter_data = 0
                for active_pulse in active_set:
                    #Currently we would not distinguish the ion_index for FPGA gates
                    temp_str = active_pulse[2].chapter_data
                    #In later experiments, the chapter data of FPGA gates for each ion is bound to be distinct
                    #Like: ('Doppler',1),('Doppler',2),('Doppler',3),('EIT',1),('EIT',2),...,('RSB',1),...

                    temp_data = int(temp_str, 2)
                    chapter_data = chapter_data | temp_data
                chapter_str = str(bin(chapter_data))
                chapter_str = chapter_str[2:32]
                chapter_str = '0'*(24-len(chapter_str)) + chapter_str
                chapter_str = chapter_str[0:8] + ' ' + chapter_str[8:16] + ' ' + chapter_str[16:24]
                result_list.append([chapter_str, time_index[idx+1]-t])

        return result_list
#####
#####

    def clear(self, make_sure = False):
        if make_sure == False:
            print("This operation will delete the sequence in current object, if you confirm \
                    to do that please set the argument \'make_sure\' to be True.")
        else:
            self.gate_sequence = []
            #del(self.gate_time)
            self.time_stamp = [0 for i in range(self.ion_number)]
            #self.gate_list = []
            # gate_list_per_ion存下每个离子对应的gate_list，list中每个元素都是一个三元组(start_time, duration, gate)
            # 其中gate是BaseGate的一个实例
            self.gate_list_per_ion = [[] for i in range(self.ion_number)]
            self.awg_trigger = False
            self.awg_trigger_time = -1
            self.current_time = -1
            self.labelled_pulse = {}


# 此函数用于画图，但传输的参数格式不能是Sequence类本身，而应该是Sequence的成员变量gate_sequence
def plot_sequence(sequence, remap_time=True, ax: Optional[plt.Axes] = None, auto_size = False, given_size = None):
    #对于数据很多的图，如果想看清的话需要设置图片大小: plt.plot(figsize=(n,m))
    #autosize == True时图片会自动设置长度，用于显示很长的sequence

    '''
    none_sequence = []
    sequence = copy.copy(sequence)
    for pulse_term in sequence:
        if pulse_term[3] == None:
            none_sequence.append(pulse_term)
    for none_pulse_term in none_sequence:
        sequence.remove(none_pulse_term)
    '''

    moments_org2rmp, moments_rmp2org = remap_to_indices(
        moment
        for (_, start_time, duration, _, _) in sequence
        for moment in (start_time, start_time + duration)
    )

    #print(moments_org2rmp, moments_rmp2org)
    xlen = len(moments_rmp2org)
    if auto_size == True:
        plt.figure(figsize=(int(xlen*1.5),5))
    elif given_size != None:
        plt.figure(figsize=given_size)

    sequence = tuple(sequence)
    ax = plt.gca() if ax is None else ax


    line1, line2 = None, None

    ion_org2rmp, ion_rmp2org = remap_to_indices(
        ion
        for (ions, start_time, duration, _, _) in sequence
        for ion in iter_int_or_tuple(ions)
    )

    for ions, start_time, duration, pulse_type, pulse_instance in sequence:
        ions = sorted(list(iter_int_or_tuple(ions)))

        if remap_time:
            remapped_start_moment = moments_org2rmp[start_time]
            remapped_end_moment = moments_org2rmp[start_time + duration]
        else:
            remapped_start_moment = start_time
            remapped_end_moment = start_time + duration

        polygon_xy = []
        for ion in ions:
            remapped_ion = ion_org2rmp[ion]

            l = remapped_start_moment
            r = remapped_end_moment
            w = r - l
            h = 4 / 5
            cx = (l + r) / 2
            cy = remapped_ion

            rect = plt.Rectangle((l, cy - h / 2), w, h,
                facecolor='white', edgecolor='black', linewidth=3,
                zorder=2.5)
            if pulse_instance.awg_flag == True:
                rect = plt.Rectangle((l, cy - h / 2), w, h,
                facecolor='white', edgecolor='blue', linewidth=3,
                zorder=2.5)
                line2 = rect
            else:
                line1 = rect

            ax.add_patch(rect)

            text = plt.Text(cx, cy, pulse_type,
                verticalalignment='center', horizontalalignment='center',
                wrap=True, color='black',
                zorder=2.5)
            ax.add_artist(text)

            polygon_xy.append((cx, cy))

        if len(ions) > 1:
            polygon = plt.Polygon(polygon_xy, False,
                edgecolor='black', linewidth=3,
                zorder=2.4)
            ax.add_patch(polygon)

    ax.set_xlabel("time")

    if remap_time:
        padding_x = 3 / 5
        ax.set_xlim(0 - padding_x, len(moments_rmp2org) - 1 + padding_x)
    else:
        padding_x = (max(moments_rmp2org) - min(moments_rmp2org)) / 20
        ax.set_xlim(min(moments_rmp2org) - padding_x, max(moments_rmp2org) + padding_x)

    ax.grid(axis='x', color='grey', ls='--')
    if remap_time:
        ax.set_xticks(range(len(moments_rmp2org)))
    else:
        ax.set_xticks(moments_rmp2org)
    ax.set_xticklabels(moments_rmp2org)

    ax.set_ylabel("ion")

    padding_y = 3 / 5
    ax.set_ylim(len(ion_rmp2org) - 1 + padding_y, -padding_y)
    ax.grid(axis='y', color='black')
    ax.set_yticks(range(len(ion_rmp2org)))
    ax.set_yticklabels(ion_rmp2org)

    ax.set_frame_on(False)

    #ax.legend(['fpga--black','awg--blue','1','1','1','1','1','1'])
    #box = ax.get_position()
    #ax.set_position([box.x0, box.y0, box.width, box.height*0.8])
    if line1 != None and line2 != None:
        ax.legend((line1,line2),('fpga','awg'), loc='upper left', bbox_to_anchor = (0,1.15), ncol=3)
    elif line1 != None:
        ax.legend([line1],['fpga'], loc='upper left', bbox_to_anchor = (0,1.15), ncol=3)
    elif line2 != None:
        ax.legend([line2],['awg'], loc='upper left', bbox_to_anchor = (0,1.15), ncol=3)

def remap_to_indices(items, key=None):
    items_set = set(items)
    items_sorted = sorted(items_set, key=key)
    moments_dict = {t: i for i, t in enumerate(items_sorted)}

    org2rmp = moments_dict
    rmp2org = items_sorted
    return org2rmp, rmp2org

def iter_int_or_tuple(item):
    try:
        for sub_item in item:
            yield sub_item
    except TypeError:
        yield item
